import numpy as np
import image_io
from image_utils import logging


def split_image_channels(image):
    '''
    该函数用于分割图像的颜色通道
    :param image: 输入图像，numpy数组
    :return: 图像的颜色通道列表
    '''
    try:
        logging.info("开始分割图像颜色通道")
          # 1. 检测图像是否为多通道图像
        if len(image.shape) < 3 or image.shape[2] < 2:
            raise ValueError("图像通道数必须大于等于2")
        
        # 2. 使用numpy的split函数，沿着通道维度分割图像
        channels = np.split(image, image.shape[2], axis=2)
        
        # 3. 把分割后的通道图像添加到一个列表中，注意要去掉多余的维度
        channel_list = [channel.squeeze(axis=2) for channel in channels]
        
        logging.info(f"成功分割图像颜色通道")
        return channel_list
    except Exception as e:
        # print(f"在分割图像颜色通道时发生错误: {e}")
        raise ValueError(f"在分割图像颜色通道时发生错误: {e}")

def merge_image_channels(file_paths):
    '''
    该函数用于合并多个单通道图像，创建一个多通道图像
    :param file_paths: 包含单通道图像文件路径的列表
    :return: 合并后的多通道图像
    '''
    try:
        logging.info("开始合并图像通道")
          # 加载每个图像文件，然后将它们合并成一个多通道图像
        # 1. 创建一个空列表来存储加载的图像
        images = []
        
        # 2. 加载每个图像文件
        for file_path in file_paths:
            img = image_io.load_image(file_path)
            
            # 检查图像是否为单通道
            if len(img.shape) > 2 and img.shape[2] > 1:
                raise ValueError("要合并子图像必须是单通道图像，不能是多通道图像")
            
            # 确保图像是2D的（单通道）
            if len(img.shape) == 2:
                # 已经是单通道，添加到列表
                images.append(img)
            else:
                # 从3D数组中提取单通道
                images.append(img[:,:,0])
        
        # 3. 检查所有图像的尺寸是否相同
        first_shape = images[0].shape
        for i, img in enumerate(images):
            if img.shape != first_shape:
                raise ValueError(f"所有图像必须具有相同的尺寸！图像 {i} 的尺寸与第一个图像不同。")
        
        # 4. 将所有单通道图像堆叠成一个多通道图像
        # 首先将每个图像扩展为3D（高度，宽度，1）
        expanded_images = [img[:,:,np.newaxis] if len(img.shape) == 2 else img for img in images]
        # 然后沿着第三维（通道维度）连接它们
        merged_image = np.concatenate(expanded_images, axis=2)
        
        logging.info("成功合并图像通道")
        return merged_image
    except Exception as e:
        raise ValueError(f"在合并图像通道时发生错误: {e}")